#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
TinyPNG Pro 服务管理脚本
用于启动和关闭前后端服务
"""

import os
import sys
import time
import signal
import subprocess
import argparse
import platform
import webbrowser
from datetime import datetime

# 配置
FRONTEND_PORT = 3000
BACKEND_PORT = 5001
FRONTEND_DIR = "frontend"
BACKEND_DIR = "backend"
LOG_DIR = os.path.join(BACKEND_DIR, "logs")

# 确保日志目录存在
os.makedirs(LOG_DIR, exist_ok=True)

# 进程记录
processes = {
    "frontend": None,
    "backend": None
}

# 日志文件
log_file = os.path.join(LOG_DIR, f"service_{datetime.now().strftime('%Y%m%d')}.log")

def log(message):
    """记录日志"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_message = f"[{timestamp}] {message}"
    print(log_message)
    
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(log_message + "\n")

def kill_process_using_port(port):
    """终止占用指定端口的进程"""
    import psutil
    for proc in psutil.process_iter(['pid', 'name', 'connections']):
        try:
            if not proc.name().startswith('python'):
                continue
            for conn in proc.connections():
                if conn.status == 'LISTEN' and conn.laddr.port == port:
                    log(f"发现端口 {port} 被进程 {proc.pid} ({proc.name()}) 占用，正在终止...")
                    proc.kill()
                    log(f"端口 {port} 被进程 {proc.pid} ({proc.name()}) 占用，终止成功")
                    return True
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            log(f"端口 {port} 被进程 {proc.pid} ({proc.name()}) 占用，终止失败")
            continue
    return False

def is_port_in_use(port):
    """检查端口是否被占用"""
    import socket
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) == 0

def start_backend():
    """启动后端服务"""
    if is_port_in_use(BACKEND_PORT):
        if not kill_process_using_port(BACKEND_PORT):
            log(f"错误: 端口 {BACKEND_PORT} 已被占用且无法终止进程")
            return False
        time.sleep(5)  # 等待进程终止
    
    log("正在启动后端服务...")
    
    # 检查Python环境
    try:
        # 安装依赖
        log("正在安装后端依赖...")
        pip_cmd = [sys.executable, "-m", "pip", "install", "-r", os.path.join(BACKEND_DIR, "requirements.txt")]
        subprocess.run(pip_cmd, check=True)
        
        # 启动后端服务
        backend_cmd = [sys.executable, os.path.join(BACKEND_DIR, "app.py")]
        
        # 根据操作系统决定如何处理输出
        if platform.system() == "Windows":
            backend_process = subprocess.Popen(
                backend_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
            )
        else:
            backend_process = subprocess.Popen(
                backend_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                preexec_fn=os.setsid
            )
        
        processes["backend"] = backend_process
        
        # 等待服务启动
        time.sleep(2)
        
        if backend_process.poll() is None:
            log(f"后端服务已启动，运行在 http://localhost:{BACKEND_PORT}")
            return True
        else:
            stdout, stderr = backend_process.communicate(timeout=1)
            log(f"后端服务启动失败: {stderr}")
            return False
            
    except Exception as e:
        log(f"启动后端服务时出错: {str(e)}")
        return False

def start_frontend():
    """启动前端服务"""
    if is_port_in_use(FRONTEND_PORT):
        if not kill_process_using_port(FRONTEND_PORT):
            log(f"错误: 端口 {FRONTEND_PORT} 已被占用且无法终止进程")
            return False
        time.sleep(1)  # 等待进程终止
    
    log("正在启动前端服务...")
    
    try:
        # 检查Node.js环境
        npm_cmd = "npm.cmd" if platform.system() == "Windows" else "npm"
        
        # 安装依赖
        log("正在安装前端依赖...")
        install_cmd = [npm_cmd, "install"]
        subprocess.run(install_cmd, cwd=FRONTEND_DIR, check=True)
        
        # 启动前端服务
        frontend_cmd = [npm_cmd, "run", "dev"]
        
        # 根据操作系统决定如何处理输出
        if platform.system() == "Windows":
            frontend_process = subprocess.Popen(
                frontend_cmd,
                cwd=FRONTEND_DIR,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
            )
        else:
            frontend_process = subprocess.Popen(
                frontend_cmd,
                cwd=FRONTEND_DIR,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                preexec_fn=os.setsid
            )
        
        processes["frontend"] = frontend_process
        
        # 等待服务启动
        time.sleep(5)
        
        if frontend_process.poll() is None:
            log(f"前端服务已启动，运行在 http://localhost:{FRONTEND_PORT}")
            return True
        else:
            stdout, stderr = frontend_process.communicate(timeout=1)
            log(f"前端服务启动失败: {stderr}")
            return False
            
    except Exception as e:
        log(f"启动前端服务时出错: {str(e)}")
        return False

def stop_service(service_name):
    """停止指定服务"""
    process = processes.get(service_name)
    if process is None or process.poll() is not None:
        log(f"{service_name}服务未运行")
        return
    
    log(f"正在停止{service_name}服务...")
    
    try:
        if platform.system() == "Windows":
            process.send_signal(signal.CTRL_BREAK_EVENT)
        else:
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
        
        # 等待进程结束
        try:
            process.wait(timeout=5)
            log(f"{service_name}服务已停止")
        except subprocess.TimeoutExpired:
            log(f"{service_name}服务未响应，强制终止")
            if platform.system() == "Windows":
                process.kill()
            else:
                os.killpg(os.getpgid(process.pid), signal.SIGKILL)
    except Exception as e:
        log(f"停止{service_name}服务时出错: {str(e)}")
        if process.poll() is None:
            process.kill()

def start_all():
    """启动所有服务"""
    backend_success = start_backend()
    frontend_success = start_frontend()
    
    if backend_success and frontend_success:
        log("所有服务已成功启动")
        # 自动打开浏览器
        webbrowser.open(f"http://localhost:{FRONTEND_PORT}")
        return True
    else:
        log("部分服务启动失败，正在清理...")
        stop_all()
        return False

def stop_all():
    """停止所有服务"""
    stop_service("frontend")
    stop_service("backend")
    log("所有服务已停止")

def check_status():
    """检查服务状态"""
    backend_running = processes["backend"] is not None and processes["backend"].poll() is None
    frontend_running = processes["frontend"] is not None and processes["frontend"].poll() is None
    
    log("服务状态:")
    log(f"- 后端服务: {'运行中' if backend_running else '已停止'}")
    log(f"- 前端服务: {'运行中' if frontend_running else '已停止'}")
    
    return backend_running, frontend_running

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="TinyPNG Pro 服务管理工具")
    parser.add_argument("action", choices=["start", "stop", "restart", "status"], help="要执行的操作")
    parser.add_argument("--backend-only", action="store_true", help="仅操作后端服务")
    parser.add_argument("--frontend-only", action="store_true", help="仅操作前端服务")
    
    args = parser.parse_args()
    
    if args.action == "start":
        if args.backend_only:
            start_backend()
        elif args.frontend_only:
            start_frontend()
        else:
            start_all()
    
    elif args.action == "stop":
        if args.backend_only:
            stop_service("backend")
        elif args.frontend_only:
            stop_service("frontend")
        else:
            stop_all()
    
    elif args.action == "restart":
        log("正在重启服务...")
        if args.backend_only:
            stop_service("backend")
            time.sleep(2)
            start_backend()
        elif args.frontend_only:
            stop_service("frontend")
            time.sleep(2)
            start_frontend()
        else:
            stop_all()
            time.sleep(2)
            start_all()
    
    elif args.action == "status":
        check_status()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        log("收到中断信号，正在停止所有服务...")
        stop_all()
    except Exception as e:
        log(f"发生错误: {str(e)}")
        stop_all()
        sys.exit(1)